#include "stdafx.h"
#include "PageHack.h"

NTSTATUS
InitMapPage(
	OUT PMAP_STRUCT MapHandle
	)
{
	NTSTATUS Status = STATUS_SUCCESS;
	PMMPTE pte;

	RtlZeroMemory(MapHandle, sizeof(*MapHandle));

	__try
	{

		MapHandle->OrigPage = ExAllocatePool(NonPagedPool,
			PAGE_SIZE*4);

		if (MapHandle->OrigPage == NULL)
		{
			Status = STATUS_INSUFFICIENT_RESOURCES;
			__leave;
		}
		MapHandle->RealPage = MapHandle->OrigPage;
		MapHandle->OrigPage = (PVOID)(((ULONG_PTR)MapHandle->RealPage + 0xFFF) & ~0xFFF);

		MapHandle->Mdl = IoAllocateMdl(MapHandle->OrigPage,
			PAGE_SIZE,
			FALSE,
			FALSE,
			NULL);

		if (MapHandle->Mdl == NULL)
		{
			Status = STATUS_INSUFFICIENT_RESOURCES;
			__leave;
		}

		//
		// Remap
		//

		MapHandle->MapPage = MmMapLockedPagesSpecifyCache(MapHandle->Mdl,
			KernelMode,
			MmCached,
			NULL,
			FALSE,
			HighPagePriority);

		if (MapHandle->MapPage == NULL)
		{
			Status = STATUS_INSUFFICIENT_RESOURCES;
			__leave;
		}

		pte = MiGetPteAddress(MapHandle->MapPage);

		MapHandle->MapPagePhys.QuadPart = *(PULONGLONG)pte;

	}
	__finally
	{

		if (!NT_SUCCESS(Status))
		{
			if (MapHandle->Mdl != NULL)
			{
				IoFreeMdl(MapHandle->Mdl);
			}

			if (MapHandle->OrigPage != NULL)
			{
				ExFreePool(MapHandle->OrigPage);
			}
		}
	}

	return Status;
}

PVOID
MapSpecifiedPage(
	IN PMAP_STRUCT MapHandle,
	IN PHYSICAL_ADDRESS PhysicalAddress
	)
{
	PMMPTE pte = MiGetPteAddress(MapHandle->MapPage);

	pte->PageFrameNumber = PhysicalAddress.QuadPart >> 12;

	_ReadWriteBarrier();

	__invlpg(MapHandle->MapPage);

	return MapHandle->MapPage;
}

VOID
FiniMapPage(
	IN PMAP_STRUCT MapHandle
	)
{
	PMMPTE pte = MiGetPteAddress(MapHandle->MapPage);

	pte->PageFrameNumber = MapHandle->MapPagePhys.QuadPart >> 12;

	MmUnmapLockedPages(MapHandle->MapPage, MapHandle->Mdl);

	IoFreeMdl(MapHandle->Mdl);

	ExFreePool(MapHandle->RealPage);
}